CycleGAN, Image-to-Image Translation

In this notebook, we're going to define and train a CycleGAN to read in an image from a set $X$ and transform it so that it looks as if it belongs in set $Y$. Specifically, we'll look at a set of images of Yosemite national park taken either during the summer of winter. The seasons are our two domains!

The objective will be to train generators that learn to transform an image from domain $X$ into an image that looks like it came from domain $Y$ (and vice versa).

Some examples of image data in both sets are pictured below.

Unpaired Training Data

These images do not come with labels, but CycleGANs give us a way to learn the mapping between one image domain and another using an unsupervised approach. A CycleGAN is designed for image-to-image translation and it learns from unpaired training data. This means that in order to train a generator to translate images from domain $X$ to domain $Y$, we do not have to have exact correspondences between individual images in those domains. For example, in the paper that introduced CycleGANs, the authors are able to translate between images of horses and zebras, even though there are no images of a zebra in exactly the same position as a horse or with exactly the same background, etc. Thus, CycleGANs enable learning a mapping from one domain $X$ to another domain $Y$ without having to find perfectly-matched, training pairs!

CycleGAN and Notebook Structure

A CycleGAN is made of two types of networks: discriminators, and generators. In this example, the discriminators are responsible for classifying images as real or fake (for both $X$ and $Y$ kinds of images). The generators are responsible for generating convincing, fake images for both kinds of images.

This notebook will detail the steps you should take to define and train such a CycleGAN.

  1. You'll load in the image data using PyTorch's DataLoader class to efficiently read in images from a specified directory.
  2. Then, you'll be tasked with defining the CycleGAN architecture according to provided specifications. You'll define the discriminator and the generator models.
  3. You'll complete the training cycle by calculating the adversarial and cycle consistency losses for the generator and discriminator network and completing a number of training epochs. It's suggested that you enable GPU usage for training.
  4. Finally, you'll evaluate your model by looking at the loss over time and looking at sample, generated images.

Load and Visualize the Data

We'll first load in and visualize the training data, importing the necessary libraries to do so.

DataLoaders

The get_data_loader function returns training and test DataLoaders that can load data efficiently and in specified batches. The function has the following parameters:

The test data is strictly for feeding to our generators, later on, so we can visualize some generated samples on fixed, test data.

You can see that this function is also responsible for making sure our images are of the right, square size (128x128x3) and converted into Tensor image types.

It's suggested that you use the default values of these parameters.

Note: If you are trying this code on a different set of data, you may get better results with larger image_size and batch_size parameters. If you change the batch_size, make sure that you create complete batches in the training loop otherwise you may get an error when trying to save sample data.

Display some Training Images

Below we provide a function imshow that reshape some given images and converts them to NumPy images so that they can be displayed by plt. This cell should display a grid that contains a batch of image data from set $X$.

Next, let's visualize a batch of images from set $Y$.

Pre-processing: scaling from -1 to 1

We need to do a bit of pre-processing; we know that the output of our tanh activated generator will contain pixel values in a range from -1 to 1, and so, we need to rescale our training images to a range of -1 to 1. (Right now, they are in a range from 0-1.)

Define the Model

A CycleGAN is made of two discriminator and two generator networks.

Discriminators

The discriminators, $D_X$ and $D_Y$, in this CycleGAN are convolutional neural networks that see an image and attempt to classify it as real or fake. In this case, real is indicated by an output close to 1 and fake as close to 0. The discriminators have the following architecture:

This network sees a 128x128x3 image, and passes it through 5 convolutional layers that downsample the image by a factor of 2. The first four convolutional layers have a BatchNorm and ReLu activation function applied to their output, and the last acts as a classification layer that outputs one value.

Convolutional Helper Function

To define the discriminators, you're expected to use the provided conv function, which creates a convolutional layer + an optional batch norm layer.

Define the Discriminator Architecture

Your task is to fill in the __init__ function with the specified 5 layer conv net architecture. Both $D_X$ and $D_Y$ have the same architecture, so we only need to define one class, and later instantiate two discriminators.

It's recommended that you use a kernel size of 4x4 and use that to determine the correct stride and padding size for each layer. This Stanford resource may also help in determining stride and padding sizes.

The forward function defines how an input image moves through the discriminator, and the most important thing is to pass it through your convolutional layers in order, with a ReLu activation function applied to all but the last layer.

You should not apply a sigmoid activation function to the output, here, and that is because we are planning on using a squared error loss for training. And you can read more about this loss function, later in the notebook.

Generators

The generators, G_XtoY and G_YtoX (sometimes called F), are made of an encoder, a conv net that is responsible for turning an image into a smaller feature representation, and a decoder, a transpose_conv net that is responsible for turning that representation into an transformed image. These generators, one from XtoY and one from YtoX, have the following architecture:

This network sees a 128x128x3 image, compresses it into a feature representation as it goes through three convolutional layers and reaches a series of residual blocks. It goes through a few (typically 6 or more) of these residual blocks, then it goes through three transpose convolutional layers (sometimes called de-conv layers) which upsample the output of the resnet blocks and create a new image!

Note that most of the convolutional and transpose-convolutional layers have BatchNorm and ReLu functions applied to their outputs with the exception of the final transpose convolutional layer, which has a tanh activation function applied to the output. Also, the residual blocks are made of convolutional and batch normalization layers, which we'll go over in more detail, next.

Residual Block Class

To define the generators, you're expected to define a ResidualBlock class which will help you connect the encoder and decoder portions of the generators. You might be wondering, what exactly is a Resnet block? It may sound familiar from something like ResNet50 for image classification, pictured below.

ResNet blocks rely on connecting the output of one layer with the input of an earlier layer. The motivation for this structure is as follows: very deep neural networks can be difficult to train. Deeper networks are more likely to have vanishing or exploding gradients and, therefore, have trouble reaching convergence; batch normalization helps with this a bit. However, during training, we often see that deep networks respond with a kind of training degradation. Essentially, the training accuracy stops improving and gets saturated at some point during training. In the worst cases, deep models would see their training accuracy actually worsen over time!

One solution to this problem is to use Resnet blocks that allow us to learn so-called residual functions as they are applied to layer inputs. You can read more about this proposed architecture in the paper, Deep Residual Learning for Image Recognition by Kaiming He et. al, and the below image is from that paper.

Residual Functions

Usually, when we create a deep learning model, the model (several layers with activations applied) is responsible for learning a mapping, M, from an input x to an output y.

M(x) = y (Equation 1)

Instead of learning a direct mapping from x to y, we can instead define a residual function

F(x) = M(x) - x

This looks at the difference between a mapping applied to x and the original input, x. F(x) is, typically, two convolutional layers + normalization layer and a ReLu in between. These convolutional layers should have the same number of inputs as outputs. This mapping can then be written as the following; a function of the residual function and the input x. The addition step creates a kind of loop that connects the input x to the output, y:

M(x) = F(x) + x (Equation 2) or

y = F(x) + x (Equation 3)

Optimizing a Residual Function

The idea is that it is easier to optimize this residual function F(x) than it is to optimize the original mapping M(x). Consider an example; what if we want y = x?

From our first, direct mapping equation, Equation 1, we could set M(x) = x but it is easier to solve the residual equation F(x) = 0, which, when plugged in to Equation 3, yields y = x.

Defining the ResidualBlock Class

To define the ResidualBlock class, we'll define residual functions (a series of layers), apply them to an input x and add them to that same input. This is defined just like any other neural network, with an __init__ function and the addition step in the forward function.

In our case, you'll want to define the residual block as:

Then, in the forward function, add the input x to this residual block. Feel free to use the helper conv function from above to create this block.

Transpose Convolutional Helper Function

To define the generators, you're expected to use the above conv function, ResidualBlock class, and the below deconv helper function, which creates a transpose convolutional layer + an optional batchnorm layer.

Define the Generator Architecture

Both $G_{XtoY}$ and $G_{YtoX}$ have the same architecture, so we only need to define one class, and later instantiate two generators.

Create the complete network

Using the classes you defined earlier, you can define the discriminators and generators necessary to create a complete CycleGAN. The given parameters should work for training.

First, create two discriminators, one for checking if $X$ sample images are real, and one for checking if $Y$ sample images are real. Then the generators. Instantiate two of them, one for transforming a painting into a realistic photo and one for transforming a photo into into a painting.

Check that you've implemented this correctly

The function create_model should return the two generator and two discriminator networks. After you've defined these discriminator and generator components, it's good practice to check your work. The easiest way to do this is to print out your model architecture and read through it to make sure the parameters are what you expected. The next cell will print out their architectures.

Discriminator and Generator Losses

Computing the discriminator and the generator losses are key to getting a CycleGAN to train.

Image from original paper by Jun-Yan Zhu et. al.

Least Squares GANs

We've seen that regular GANs treat the discriminator as a classifier with the sigmoid cross entropy loss function. However, this loss function may lead to the vanishing gradients problem during the learning process. To overcome such a problem, we'll use a least squares loss function for the discriminator. This structure is also referred to as a least squares GAN or LSGAN, and you can read the original paper on LSGANs, here. The authors show that LSGANs are able to generate higher quality images than regular GANs and that this loss type is a bit more stable during training!

Discriminator Losses

The discriminator losses will be mean squared errors between the output of the discriminator, given an image, and the target value, 0 or 1, depending on whether it should classify that image as fake or real. For example, for a real image, x, we can train $D_X$ by looking at how close it is to recognizing and image x as real using the mean squared error:

out_x = D_X(x)
real_err = torch.mean((out_x-1)**2)

Generator Losses

Calculating the generator losses will look somewhat similar to calculating the discriminator loss; there will still be steps in which you generate fake images that look like they belong to the set of $X$ images but are based on real images in set $Y$, and vice versa. You'll compute the "real loss" on those generated images by looking at the output of the discriminator as it's applied to these fake images; this time, your generator aims to make the discriminator classify these fake images as real images.

Cycle Consistency Loss

In addition to the adversarial losses, the generator loss terms will also include the cycle consistency loss. This loss is a measure of how good a reconstructed image is, when compared to an original image.

Say you have a fake, generated image, x_hat, and a real image, y. You can get a reconstructed y_hat by applying G_XtoY(x_hat) = y_hat and then check to see if this reconstruction y_hat and the orginal image y match. For this, we recommed calculating the L1 loss, which is an absolute difference, between reconstructed and real images. You may also choose to multiply this loss by some weight value lambda_weight to convey its importance.

The total generator loss will be the sum of the generator losses and the forward and backward cycle consistency losses.

Define Loss Functions

To help us calculate the discriminator and gnerator losses during training, let's define some helpful loss functions. Here, we'll define three.

  1. real_mse_loss that looks at the output of a discriminator and returns the error based on how close that output is to being classified as real. This should be a mean squared error.
  2. fake_mse_loss that looks at the output of a discriminator and returns the error based on how close that output is to being classified as fake. This should be a mean squared error.
  3. cycle_consistency_loss that looks at a set of real image and a set of reconstructed/generated images, and returns the mean absolute error between them. This has a lambda_weight parameter that will weight the mean absolute error in a batch.

It's recommended that you take a look at the original, CycleGAN paper to get a starting value for lambda_weight.

Define the Optimizers

Next, let's define how this model will update its weights. This, like the GANs you may have seen before, uses Adam optimizers for the discriminator and generator. It's again recommended that you take a look at the original, CycleGAN paper to get starting hyperparameter values.

Training a CycleGAN

When a CycleGAN trains, and sees one batch of real images from set $X$ and $Y$, it trains by performing the following steps:

Training the Discriminators

  1. Compute the discriminator $D_X$ loss on real images
  2. Generate fake images that look like domain $X$ based on real images in domain $Y$
  3. Compute the fake loss for $D_X$
  4. Compute the total loss and perform backpropagation and $D_X$ optimization
  5. Repeat steps 1-4 only with $D_Y$ and your domains switched!

Training the Generators

  1. Generate fake images that look like domain $X$ based on real images in domain $Y$
  2. Compute the generator loss based on how $D_X$ responds to fake $X$
  3. Generate reconstructed $\hat{Y}$ images based on the fake $X$ images generated in step 1
  4. Compute the cycle consistency loss by comparing the reconstructions with real $Y$ images
  5. Repeat steps 1-4 only swapping domains
  6. Add up all the generator and reconstruction losses and perform backpropagation + optimization

Saving Your Progress

A CycleGAN repeats its training process, alternating between training the discriminators and the generators, for a specified number of training iterations. You've been given code that will save some example generated images that the CycleGAN has learned to generate after a certain number of training iterations. Along with looking at the losses, these example generations should give you an idea of how well your network has trained.

Below, you may choose to keep all default parameters; your only task is to calculate the appropriate losses and complete the training cycle.

Tips on Training and Loss Patterns

A lot of experimentation goes into finding the best hyperparameters such that the generators and discriminators don't overpower each other. It's often a good starting point to look at existing papers to find what has worked in previous experiments, I'd recommend this DCGAN paper in addition to the original CycleGAN paper to see what worked for them. Then, you can try your own experiments based off of a good foundation.

Discriminator Losses

When you display the generator and discriminator losses you should see that there is always some discriminator loss; recall that we are trying to design a model that can generate good "fake" images. So, the ideal discriminator will not be able to tell the difference between real and fake images and, as such, will always have some loss. You should also see that $D_X$ and $D_Y$ are roughly at the same loss levels; if they are not, this indicates that your training is favoring one type of discriminator over the and you may need to look at biases in your models or data.

Generator Loss

The generator's loss should start significantly higher than the discriminator losses because it is accounting for the loss of both generators and weighted reconstruction errors. You should see this loss decrease a lot at the start of training because initial, generated images are often far-off from being good fakes. After some time it may level off; this is normal since the generator and discriminator are both improving as they train. If you see that the loss is jumping around a lot, over time, you may want to try decreasing your learning rates or changing your cycle consistency loss to be a little more/less weighted.

Evaluate the Result!

As you trained this model, you may have chosen to sample and save the results of your generated images after a certain number of training iterations. This gives you a way to see whether or not your Generators are creating good fake images. For example, the image below depicts real images in the $Y$ set, and the corresponding generated images during different points in the training process. You can see that the generator starts out creating very noisy, fake images, but begins to converge to better representations as it trains (though, not perfect).

Below, you've been given a helper function for displaying generated samples based on the passed in training iteration.

Further Challenges and Directions

Once you are satified with your model, you are ancouraged to test it on a different dataset to see if it can find different types of mappings!

Different datasets for download

You can download a variety of datasets used in the Pix2Pix and CycleGAN papers, by following instructions in the associated Github repository. You'll just need to make sure that the data directories are named and organized correctly to load in that data.